"""
该类用于数据的提取
"""
import copy
import logging
import traceback
from io import BytesIO

import torch.utils.data as data
import os.path
from numpy.random import randint
from torch.utils.data import DataLoader

from tools.transforms import GroupMultiScaleCrop, GroupRandomHorizontalFlip, Stack, ToTorchFormatTensor, \
    IdentityTransform

import os
from tools.transforms import *


class VideoRecord(object):
    def __init__(self,path_prefix, row):
        row[0] = os.path.join(path_prefix,row[0])
        self._data = row

    # 放置的图片位置
    @property
    def path(self):
        return self._data[0]

    # 帧数数量
    @property
    def num_frames(self):
        return int(self._data[1])

    # 所属标签
    @property
    def label(self):
        return int(self._data[2])


class UcfFramesDataSet(data.Dataset):
    def __init__(self,
                 ann_file,
                 t_num,
                 image_tmpl="img_{:05d}.jpg",
                 transform=None,
                 random_shift=True,
                 test_mode=False,
                 log= logging.getLogger("root"),
                 root_dir = "",
                 ):
        self.ann_file = ann_file
        self.t_num = t_num
        self.image_tmpl = image_tmpl
        self.transform = transform
        self.random_shift = random_shift
        self.test_mode = test_mode
        self.log = log

        self.video_list = [VideoRecord(root_dir,x.strip().split(' ')) for x in open(self.ann_file)]

    def __getitem__(self, index):
        return self.get(index)

    def get(self, index):
        record = self.video_list[index]
        if not self.test_mode:
            segment_indices = self._sample_indices(record) if self.random_shift else self._get_val_indices(record)
        else:
            segment_indices = self._get_test_indices(record)

        images = []
        for p in segment_indices:
            images.append(self.transform(self._load_image(record.path, p)))
        # 转变为Tensor
        data = torch.stack(images,0)
        return data, record.label


    def __len__(self):
        return len(self.video_list)

    def _load_image(self, directory, idx):
        # 这里可能会遇到被损坏的图片
        path = os.path.join(directory, self.image_tmpl.format(idx))
        try:
            img = Image.open(path)
            return img.convert('RGB')
        except Exception as e:
            self.log.error("ERROR Path %s" % path)
            self.log.error(traceback.print_exc())
            return None

    def _get_test_indices(self, record):
        tick = record.num_frames / float(self.t_num)
        offsets = np.array([int(tick / 2.0 + tick * x) for x in range(self.t_num)])
        return offsets + 1

    def _get_val_indices(self, record):
        if record.num_frames > self.t_num:
            tick = record.num_frames / float(self.t_num)
            offsets = np.array([int(tick / 2.0 + tick * x) for x in range(self.t_num)])
        else:
            raise RuntimeError
        return offsets + 1

    def _sample_indices(self, record):
        """
        :param record: VideoRecord
        :return: list
        """
        # 随机抽取几个帧的照片，数量为t_num
        average_duration = record.num_frames // self.t_num
        if average_duration > 0:
            offsets = np.multiply(list(range(self.t_num)), average_duration) + randint(average_duration,
                                                                                              size=self.t_num)
        elif record.num_frames > self.t_num:
            raise RuntimeError
            # offsets = np.sort(randint(record.num_frames - self.new_length + 1, size=self.t_num))
        else:
            raise RuntimeError
            # offsets = np.zeros((self.t_num,))
        return offsets + 1

